package com.liu.ai.store;

import io.micrometer.observation.ObservationRegistry;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.ollama.management.ModelManagementOptions;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.jdbc.core.JdbcTemplate;

import java.util.List;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

public class VectorStoreDemo {

    public static void main(String[] args) {

        VectorStore vectorStore = getVectorStore();
        List<Document> doc = vectorStore.similaritySearch("你是谁");

        Filter.Expression expression = new FilterExpressionBuilder()
                .eq("auther", "")
                .build();
        // 构建查询请求
        List<Document> documents = vectorStore.similaritySearch(SearchRequest.builder()
                .query("你是谁")
                .topK(5)
                .similarityThreshold(0.6)
                .filterExpression(expression)
                .build());
    }

    public static VectorStore getVectorStore() {

        JdbcTemplate jdbcTemplate = new JdbcTemplate(JdbcTools.ds);
        OllamaApi api = new OllamaApi();

        OllamaEmbeddingModel embedding = new OllamaEmbeddingModel(
                api,
                OllamaOptions.builder()
                        .model("zyw0605688/gte-large-zh:latest")
                        .build(),
                ObservationRegistry.create(),
                ModelManagementOptions.builder().build());
        return PgVectorStore.builder(jdbcTemplate, embedding)
                //.dimensions(1024)                    // Optional: defaults to model dimensions or 1536
                .distanceType(COSINE_DISTANCE)       // Optional: defaults to COSINE_DISTANCE
                .indexType(HNSW)                     // Optional: defaults to HNSW
                .initializeSchema(true)              // Optional: defaults to false
                .schemaName("public")                // Optional: defaults to "public"
                .vectorTableName("vector_store")     // Optional: defaults to "vector_store"
                .maxDocumentBatchSize(10000)         // Optional: defaults to 10000
                .build();
    }
}
